# RBC教材第四章可变劳动的例子

rm(list = ls())
devtools::load_all()
library(pacman)
p_load(rootSolve,signal,magrittr,ggplot2)
# 参数设置
deltax <- 0.1
thetax <- 0.36
betax <- 0.98
A <- 0.5

# 数值计算设置
crit <- 1
tol <- 0.001
eps <- 0.02

# 格点和值函数初值设置和下一期最优资本存量内存分配
cgrid <- seq(0.6,1,length.out = 80)
k0 <- seq(2,10,length.out = 80)
v <- seq(-4,9,length.out = length(k0))

bellman_sg <- function(c0, a0, agrid, vold, thetax, deltax, betax, A){

  # 通过一阶条件获得先求解劳动
  chfun <- function(h){
    # c0 <- a0^thetax*h^(1-thetax) - a1 + (1-deltax)*a0
    A*c0/(1-h)-(1-thetax)*a0^thetax*h^(-thetax)
  }
  h0 <- uniroot.all(chfun, interval = c(0,1))
  # print(h0)

  a1 <- a0^thetax*h0^(1-thetax) - c0 + (1-deltax)*a0

  if (c0 < 0) return(-1e10)
  # if (a1 > agrid[length(agrid)]) return(a1^2*(-1e10))

  return(log(c0) + A*log(1-h0) + betax * interp1(agrid, vold, xi = a1))
}

 tmp(v)$v

ans <- Sys.time()
j <- 0
while (crit > tol) {
  j <- j + 1
  print(crit)
  v <- tmap(v)$v
  crit <- abs(vold-v) %>% mean() %>% max()
}
Sys.time()- ans

# 获得劳动力
chfun <- function(h,a0,a1,A,thetax,deltax){
  c0 <- a0^thetax*h^(1-thetax) - a1 + (1-deltax)*a0
  A*c0/(1-h)-(1-thetax)*a0^thetax*h^(-thetax)
}

hh <- numeric(80)
for (i in 1:80) {
  hh[i] <- uniroot.all(chfun, interval = c(0,1), a0 = k0[i], a1 = aopt[i], A = A,
                       thetax = thetax, deltax = deltax)
}
# 获得消费
c0 <- k0^thetax*hh^(1-thetax) - aopt + (1-deltax)*k0


picdata <- cbind(k0,aopt,hh,c0,v) %>% as.data.frame()
ggplot(picdata, aes(x = k0, y = aopt)) + geom_line() +
  geom_line(aes(y = k0), color = 'red')

